package com.uraurora.crescent.advanced.tree.impl;

import com.uraurora.crescent.advanced.tree.IDisjointSet;
import com.google.common.collect.Maps;
import com.uraurora.crescent.tuple.Tuple2;

import java.util.*;
import java.util.stream.Collectors;

public class DisjointSet<T> implements IDisjointSet<T> {

    private int size;

    private Map<T, InternalNode<T>> map;

    public DisjointSet(){
        this.size = 0;
        this.map = Maps.newHashMap();
    }

    @Override
    public boolean isConnected(T a, T b) {
        return find(a).equals(find(b));
    }

    @Override
    public void union(T a, T b) {
        if (!map.containsKey(a)) {
            InternalNode<T> val = new InternalNode<>(a, 0);
            val.parent = val;
            map.put(a, val);
            size++;
        }
        if(!map.containsKey(b)){
            InternalNode<T> val = new InternalNode<>(b, 0);
            val.parent = val;
            map.put(b, val);
            size++;
        }
        InternalNode<T> pRoot = find(a), qRoot = find(b);

        if( pRoot.equals(qRoot)) {
            return;
        }
        if( pRoot.rank < qRoot.rank ){
            pRoot.parent = qRoot;
        }
        else if( qRoot.rank < pRoot.rank){
            qRoot.parent = pRoot;
        }
        else{
            pRoot.parent = qRoot;
            qRoot.rank ++;
        }
    }

    @SuppressWarnings("unchecked")
    public void union(Iterable<Tuple2<T, T>> element){
        element.forEach(item->this.union((T) item.get(0), (T) item.get(1)));
    }

    @Override
    public int size() {
        return this.size;
    }

    @Override
    public int unionSize() {
        return map.keySet().stream().map(this::find).collect(Collectors.toSet()).size();
    }

    static class InternalNode<T>{
        private final T val;
        private InternalNode<T> parent;
        private int rank;

        InternalNode(T val, int rank){
            this.val = val;
            this.rank = rank;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            InternalNode<?> that = (InternalNode<?>) o;
            return Objects.equals(val, that.val);
        }

        @Override
        public int hashCode() {
            return Objects.hash(val);
        }
    }

    private InternalNode<T>find(T node){
        assert map.containsKey(node);
        InternalNode<T>n = map.get(node);
        return n.parent.equals(n) ? n : find(n.parent.val);
    }
}
